"""
Redis连接管理模块
用于将处理结果存入Redis
"""

import redis
import logging
import json
import time
import base64
from typing import Dict, Any, Optional

# 配置日志
logger = logging.getLogger(__name__)

class RedisManager:
    """Redis管理器，用于连接Redis并存储处理结果"""
    
    def __init__(self, host="localhost", port=6379, password="", db=0):
        """
        初始化Redis连接
        
        Args:
            host: Redis主机名
            port: Redis端口
            password: Redis密码
            db: Redis数据库索引
        """
        self.host = host
        self.port = port
        self.password = password
        self.db = db
        self.client = None
        
        # 尝试连接Redis
        self.connect()
    
    def connect(self) -> bool:
        """
        连接到Redis
        
        Returns:
            连接是否成功
        """
        try:
            # 如果密码为空，不传递password参数
            if self.password:
                self.client = redis.Redis(
                    host=self.host,
                    port=self.port,
                    password=self.password,
                    db=self.db,
                    decode_responses=False  # 保持字节格式以支持二进制数据
                )
            else:
                self.client = redis.Redis(
                    host=self.host,
                    port=self.port,
                    db=self.db,
                    decode_responses=False  # 保持字节格式以支持二进制数据
                )
            # 测试连接
            self.client.ping()
            logger.info(f"成功连接到Redis {self.host}:{self.port}")
            return True
        except Exception as e:
            logger.error(f"连接Redis失败: {e}")
            self.client = None
            return False
    
    def reconnect_if_needed(self) -> bool:
        """
        如果连接断开，重新连接
        
        Returns:
            重连是否成功
        """
        try:
            if self.client:
                self.client.ping()
                return True
            else:
                logger.info("Redis连接已断开，尝试重新连接")
                return self.connect()
        except Exception as e:
            logger.warning(f"检查Redis连接时出错: {e}")
            return self.connect()
    
    def save_detection_result(self, redis_key: str, data: Dict[str, Any]) -> bool:
        """
        保存检测结果到Redis
        
        Args:
            redis_key: Redis键名
            data: 要保存的数据，包含以下字段:
                - person_count: 人数
                - helmet_count: 安全帽数量
                - gap_person_and_helmet: 未佩戴安全帽的人数
                - image_base64: 图像的Base64编码字符串
                - frame_count: 帧计数
                
        Returns:
            保存是否成功
        """
        if not self.reconnect_if_needed():
            logger.error("无法连接到Redis，无法保存检测结果")
            return False
        
        try:
            # 保存主要结果到hash - 所有值都转换为字符串避免序列化问题
            if 'person_count' in data:
                self.client.hset(redis_key, 'person_count', str(data['person_count']))
            
            if 'helmet_count' in data:
                self.client.hset(redis_key, 'helmet_count', str(data['helmet_count']))
            
            if 'gap_person_and_helmet' in data:
                self.client.hset(redis_key, 'gap_person_and_helmet', str(data['gap_person_and_helmet']))
            
            # 保存图像数据（如果有）
            if 'image_base64' in data:
                # 直接保存base64字符串，避免Java端序列化问题
                self.client.hset(redis_key, 'img_byte', data['image_base64'])
            
            # 保存时间戳 - 转换为字符串
            self.client.hset(redis_key, 'update_time', str(int(time.time() * 1000)))
            
            # 保存帧号（如果有）- 转换为字符串
            if 'frame_count' in data:
                self.client.hset(redis_key, 'frame_count', str(data['frame_count']))
            
            # 设置过期时间（24小时）
            self.client.expire(redis_key, 86400)
            
            logger.debug(f"数据已保存到Redis，键={redis_key}")
            return True
            
        except Exception as e:
            logger.error(f"保存到Redis失败: {e}")
            return False
    
    def get_detection_result(self, redis_key: str) -> Optional[Dict[str, Any]]:
        """
        从Redis获取检测结果
        
        Args:
            redis_key: Redis键名
            
        Returns:
            检测结果或None
        """
        if not self.reconnect_if_needed():
            logger.error("无法连接到Redis，无法获取检测结果")
            return None
        
        try:
            # 获取所有字段
            data = self.client.hgetall(redis_key)
            
            if not data:
                return None
            
            # 转换为字典
            result = {}
            for key, value in data.items():
                key_str = key.decode('utf-8') if isinstance(key, bytes) else key
                
                # 处理特殊字段
                if key_str == 'img_byte':
                    # 图像数据保持原样
                    result[key_str] = value
                else:
                    # 其他字段转换为字符串
                    value_str = value.decode('utf-8') if isinstance(value, bytes) else value
                    result[key_str] = value_str
            
            return result
            
        except Exception as e:
            logger.error(f"从Redis获取数据失败: {e}")
            return None
    
    def close(self):
        """关闭Redis连接"""
        if self.client:
            self.client = None
            logger.info("Redis连接已关闭")
    
    def __del__(self):
        """析构函数，确保关闭连接"""
        self.close()


if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(level=logging.INFO)
    
    # 测试Redis连接
    redis_manager = RedisManager()
    
    # 测试保存数据
    test_data = {
        'person_count': 5,
        'helmet_count': 3,
        'gap_person_and_helmet': 2,
        'frame_count': 100,
        'image_base64': 'test_base64_data'
    }
    
    success = redis_manager.save_detection_result('test_key', test_data)
    print(f"保存结果: {success}")
    
    # 测试获取数据
    result = redis_manager.get_detection_result('test_key')
    print(f"获取结果: {result}")